{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021-2023 @ Shenzhen Bay Laboratory & Peking University & Huawei Technologies Co., Ltd\n",
    "\n",
    "This code is a part of MindSPONGE:\n",
    "MindSpore Simulation Package tOwards Next Generation molecular modelling.\n",
    "\n",
    "MindSPONGE is open-source software based on the AI-framework:\n",
    "MindSpore (https://www.mindspore.cn/)\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "\n",
    "You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "\n",
    "MindSPONGE advanced tutorial 03: Energy wrapper and integrated tempering sampling (ITS)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindspore import context\n",
    "from mindsponge import Sponge\n",
    "from mindsponge import ForceField\n",
    "from mindsponge import set_global_units\n",
    "from mindsponge import Protein\n",
    "from mindsponge import WithEnergyCell\n",
    "from mindsponge import UpdaterMD\n",
    "from mindsponge.optimizer import SteepestDescent\n",
    "from mindsponge.callback import WriteH5MD, RunInfo\n",
    "from mindsponge.sampling import ITS\n",
    "from mindsponge.function import VelocityGenerator\n",
    "from mindsponge.colvar import Torsion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "context.set_context(mode=context.PYNATIVE_MODE, device_target=\"GPU\")\n",
    "set_global_units('nm', 'kj/mol')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "system = Protein(pdb='alad.pdb')\n",
    "potential = ForceField(system, 'AMBER.FF14SB')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "phi = Torsion([4, 6, 8, 14])\n",
    "psi = Torsion([6, 8, 14, 16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MindSPONGE] Using fully connected neighbour list (not updated).\n"
     ]
    }
   ],
   "source": [
    "min_opt = SteepestDescent(system.trainable_params(), 1e-7)\n",
    "mini = Sponge(system, potential, min_opt, metrics={'phi': phi, 'psi': psi})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MindSPONGE] Started simulation at 2023-05-19 10:23:42\n",
      "[MindSPONGE] Step: 0, E_pot: -55.75643, phi: -3.140886, psi: -3.1410766\n",
      "[MindSPONGE] Step: 10, E_pot: -58.40555, phi: -3.1340268, psi: -3.1368413\n",
      "[MindSPONGE] Step: 20, E_pot: -60.168476, phi: -3.1277595, psi: -3.1337605\n",
      "[MindSPONGE] Step: 30, E_pot: -61.549706, phi: -3.1221528, psi: -3.131402\n",
      "[MindSPONGE] Step: 40, E_pot: -62.696198, phi: -3.1171532, psi: -3.1295478\n",
      "[MindSPONGE] Step: 50, E_pot: -63.674088, phi: -3.112676, psi: -3.1280725\n",
      "[MindSPONGE] Step: 60, E_pot: -64.52254, phi: -3.1086385, psi: -3.1268952\n",
      "[MindSPONGE] Step: 70, E_pot: -65.26811, phi: -3.1049678, psi: -3.1259575\n",
      "[MindSPONGE] Step: 80, E_pot: -65.93031, phi: -3.1016028, psi: -3.1252148\n",
      "[MindSPONGE] Step: 90, E_pot: -66.52346, phi: -3.0984938, psi: -3.1246321\n",
      "[MindSPONGE] Finished simulation at 2023-05-19 10:23:46\n",
      "[MindSPONGE] Simulation time: 3.78 seconds.\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<mindsponge.core.sponge.Sponge at 0x7fe36945d6d0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "run_info = RunInfo(10)\n",
    "mini.run(100, callbacks=[run_info])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "its = ITS(\n",
    "    sim_temp=300,\n",
    "    temp_min=270,\n",
    "    temp_max=470,\n",
    "    temp_bin=100,\n",
    "    update_pace=100,\n",
    "    unlinear_temp=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MindSPONGE] Using fully connected neighbour list (not updated).\n"
     ]
    }
   ],
   "source": [
    "sim = WithEnergyCell(system, potential, wrapper=its)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "vgen = VelocityGenerator(300)\n",
    "velocity = vgen(system.shape, system.atom_mass)\n",
    "opt = UpdaterMD(\n",
    "    system=system,\n",
    "    time_step=1e-3,\n",
    "    velocity=velocity,\n",
    "    integrator='velocity_verlet',\n",
    "    temperature=300,\n",
    "    thermostat='langevin',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "md = Sponge(sim, optimizer=opt, metrics={'phi': phi, 'psi': psi})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MindSPONGE] Started simulation at 2023-05-19 10:23:52\n",
      "[MindSPONGE] Step: 0, E_pot: -79.28045, E_kin: 96.96472, E_tot: 17.684273, Temperature: 370.2278, phi: -3.0946958, psi: -3.1193402\n",
      "[MindSPONGE] Step: 10, E_pot: -99.59543, E_kin: 35.44615, E_tot: -64.14928, Temperature: 135.33943, phi: -3.079391, psi: -3.0574105\n",
      "[MindSPONGE] Step: 20, E_pot: -111.30186, E_kin: 49.054054, E_tot: -62.247803, Temperature: 187.29672, phi: -3.0604281, psi: -2.9846\n",
      "[MindSPONGE] Step: 30, E_pot: -108.976944, E_kin: 45.604313, E_tot: -63.37263, Temperature: 174.12502, phi: -3.1081102, psi: -2.9990194\n",
      "[MindSPONGE] Step: 40, E_pot: -119.62045, E_kin: 65.65898, E_tot: -53.96147, Temperature: 250.69716, phi: 3.1262376, psi: -3.109618\n",
      "[MindSPONGE] Step: 50, E_pot: -102.952934, E_kin: 62.735622, E_tot: -40.21731, Temperature: 239.53528, phi: -3.1144304, psi: 3.1358683\n",
      "[MindSPONGE] Step: 60, E_pot: -109.93755, E_kin: 72.00345, E_tot: -37.934105, Temperature: 274.92142, phi: -2.9780388, psi: -3.0246732\n",
      "[MindSPONGE] Step: 70, E_pot: -94.96619, E_kin: 56.707794, E_tot: -38.258392, Temperature: 216.52, phi: -2.908525, psi: -2.9314706\n",
      "[MindSPONGE] Step: 80, E_pot: -104.663734, E_kin: 83.492096, E_tot: -21.171638, Temperature: 318.78702, phi: -2.943411, psi: -2.9434223\n",
      "[MindSPONGE] Step: 90, E_pot: -97.40086, E_kin: 67.62592, E_tot: -29.77494, Temperature: 258.20728, phi: -2.920756, psi: -2.9658585\n",
      "[MindSPONGE] Step: 100, E_pot: -72.45053, E_kin: 59.22229, E_tot: -13.228241, Temperature: 226.12077, phi: -2.832687, psi: -2.9637332\n",
      "[MindSPONGE] Step: 110, E_pot: -82.195045, E_kin: 72.86748, E_tot: -9.327568, Temperature: 278.22043, phi: -2.7320318, psi: -2.9288342\n",
      "[MindSPONGE] Step: 120, E_pot: -97.2785, E_kin: 77.44091, E_tot: -19.837593, Temperature: 295.68256, phi: -2.7042341, psi: -2.9305668\n",
      "[MindSPONGE] Step: 130, E_pot: -84.327866, E_kin: 60.96147, E_tot: -23.366394, Temperature: 232.76128, phi: -2.715205, psi: -2.9548028\n",
      "[MindSPONGE] Step: 140, E_pot: -84.51745, E_kin: 60.497833, E_tot: -24.019615, Temperature: 230.99101, phi: -2.7281532, psi: -2.9871135\n",
      "[MindSPONGE] Step: 150, E_pot: -85.03114, E_kin: 64.736984, E_tot: -20.294159, Temperature: 247.17682, phi: -2.702565, psi: -2.982405\n",
      "[MindSPONGE] Step: 160, E_pot: -115.53044, E_kin: 85.95116, E_tot: -29.579285, Temperature: 328.17615, phi: -2.6703272, psi: -2.948932\n",
      "[MindSPONGE] Step: 170, E_pot: -85.69896, E_kin: 51.34121, E_tot: -34.35775, Temperature: 196.02948, phi: -2.671347, psi: -2.8957093\n",
      "[MindSPONGE] Step: 180, E_pot: -104.38485, E_kin: 64.17493, E_tot: -40.209923, Temperature: 245.03078, phi: -2.7670739, psi: -2.8358617\n",
      "[MindSPONGE] Step: 190, E_pot: -93.73041, E_kin: 60.18208, E_tot: -33.54833, Temperature: 229.78542, phi: -2.9059458, psi: -2.8860497\n",
      "[MindSPONGE] Step: 200, E_pot: -115.261795, E_kin: 66.637184, E_tot: -48.62461, Temperature: 254.4321, phi: -2.9260824, psi: -3.0448751\n",
      "[MindSPONGE] Step: 210, E_pot: -108.79744, E_kin: 52.887497, E_tot: -55.909943, Temperature: 201.93347, phi: -2.8544319, psi: -3.108961\n",
      "[MindSPONGE] Step: 220, E_pot: -110.16222, E_kin: 54.030464, E_tot: -56.131752, Temperature: 206.29752, phi: -2.823438, psi: -3.0502272\n",
      "[MindSPONGE] Step: 230, E_pot: -94.68549, E_kin: 51.776966, E_tot: -42.908527, Temperature: 197.69327, phi: -2.8482459, psi: -2.9865818\n",
      "[MindSPONGE] Step: 240, E_pot: -103.66879, E_kin: 65.93643, E_tot: -37.73236, Temperature: 251.7565, phi: -2.8394794, psi: -2.993554\n",
      "[MindSPONGE] Step: 250, E_pot: -108.18496, E_kin: 72.079895, E_tot: -36.105064, Temperature: 275.21332, phi: -2.8446143, psi: -3.091453\n",
      "[MindSPONGE] Step: 260, E_pot: -95.51515, E_kin: 56.818645, E_tot: -38.696507, Temperature: 216.94325, phi: -2.834859, psi: -3.1161938\n",
      "[MindSPONGE] Step: 270, E_pot: -96.14642, E_kin: 60.079567, E_tot: -36.066856, Temperature: 229.39401, phi: -2.8278198, psi: -3.116531\n",
      "[MindSPONGE] Step: 280, E_pot: -106.9625, E_kin: 71.73568, E_tot: -35.22682, Temperature: 273.89905, phi: -2.8131237, psi: -3.1007457\n",
      "[MindSPONGE] Step: 290, E_pot: -99.29444, E_kin: 57.426342, E_tot: -41.8681, Temperature: 219.26355, phi: -2.848291, psi: 3.1322818\n",
      "[MindSPONGE] Step: 300, E_pot: -113.620224, E_kin: 70.66248, E_tot: -42.95774, Temperature: 269.8014, phi: -2.8715844, psi: 3.1145291\n",
      "[MindSPONGE] Step: 310, E_pot: -118.13048, E_kin: 72.34643, E_tot: -45.78405, Temperature: 276.23096, phi: -2.8429792, psi: 3.0604239\n",
      "[MindSPONGE] Step: 320, E_pot: -98.75694, E_kin: 64.25189, E_tot: -34.50505, Temperature: 245.32466, phi: -2.7891164, psi: 3.0627818\n",
      "[MindSPONGE] Step: 330, E_pot: -100.50948, E_kin: 59.17732, E_tot: -41.332165, Temperature: 225.94907, phi: -2.8301227, psi: -3.1138065\n",
      "[MindSPONGE] Step: 340, E_pot: -101.83601, E_kin: 68.27479, E_tot: -33.561226, Temperature: 260.68475, phi: -2.9252996, psi: -3.1164124\n",
      "[MindSPONGE] Step: 350, E_pot: -93.47495, E_kin: 59.13736, E_tot: -34.337593, Temperature: 225.7965, phi: -2.9022496, psi: -3.1399508\n",
      "[MindSPONGE] Step: 360, E_pot: -96.12489, E_kin: 71.56705, E_tot: -24.557846, Temperature: 273.25516, phi: -2.8603244, psi: 3.127882\n",
      "[MindSPONGE] Step: 370, E_pot: -102.83362, E_kin: 92.22775, E_tot: -10.6058655, Temperature: 352.14127, phi: -2.8185446, psi: 3.1236641\n",
      "[MindSPONGE] Step: 380, E_pot: -101.22379, E_kin: 92.19026, E_tot: -9.033531, Temperature: 351.9981, phi: -2.7734127, psi: 3.139261\n",
      "[MindSPONGE] Step: 390, E_pot: -81.82995, E_kin: 68.22479, E_tot: -13.605156, Temperature: 260.49387, phi: -2.7406216, psi: 3.1272378\n",
      "[MindSPONGE] Step: 400, E_pot: -87.21591, E_kin: 73.05925, E_tot: -14.156662, Temperature: 278.95264, phi: -2.7589998, psi: 3.0983782\n",
      "[MindSPONGE] Step: 410, E_pot: -94.875725, E_kin: 70.047516, E_tot: -24.828209, Temperature: 267.4533, phi: -2.860981, psi: 3.024408\n",
      "[MindSPONGE] Step: 420, E_pot: -91.76773, E_kin: 67.19333, E_tot: -24.574402, Temperature: 256.55554, phi: -2.8950741, psi: 2.9820845\n",
      "[MindSPONGE] Step: 430, E_pot: -116.06351, E_kin: 75.71397, E_tot: -40.34954, Temperature: 289.0888, phi: -2.8017137, psi: 3.0040555\n",
      "[MindSPONGE] Step: 440, E_pot: -118.035385, E_kin: 72.956604, E_tot: -45.07878, Temperature: 278.56073, phi: -2.7482436, psi: 3.0788944\n",
      "[MindSPONGE] Step: 450, E_pot: -117.01689, E_kin: 67.02751, E_tot: -49.98938, Temperature: 255.92244, phi: -2.773432, psi: 3.104705\n",
      "[MindSPONGE] Step: 460, E_pot: -109.07965, E_kin: 60.426334, E_tot: -48.653316, Temperature: 230.71802, phi: -2.842889, psi: 3.0553405\n",
      "[MindSPONGE] Step: 470, E_pot: -106.05369, E_kin: 57.395294, E_tot: -48.658394, Temperature: 219.145, phi: -2.9089627, psi: 3.015567\n",
      "[MindSPONGE] Step: 480, E_pot: -109.92711, E_kin: 61.700203, E_tot: -48.226906, Temperature: 235.58188, phi: -2.9370558, psi: 3.0121036\n",
      "[MindSPONGE] Step: 490, E_pot: -88.65821, E_kin: 49.596222, E_tot: -39.06199, Temperature: 189.36682, phi: -2.8904948, psi: 2.9450274\n",
      "[MindSPONGE] Step: 500, E_pot: -105.0544, E_kin: 71.4137, E_tot: -33.6407, Temperature: 272.66965, phi: -2.7598214, psi: 2.8192878\n",
      "[MindSPONGE] Step: 510, E_pot: -88.247665, E_kin: 58.4207, E_tot: -29.826965, Temperature: 223.06017, phi: -2.6674225, psi: 2.7770677\n",
      "[MindSPONGE] Step: 520, E_pot: -106.27849, E_kin: 71.92129, E_tot: -34.3572, Temperature: 274.60773, phi: -2.7309775, psi: 2.7746768\n",
      "[MindSPONGE] Step: 530, E_pot: -80.48755, E_kin: 80.690735, E_tot: 0.20318604, Temperature: 308.09097, phi: -2.7848978, psi: 2.7718053\n",
      "[MindSPONGE] Step: 540, E_pot: -100.28407, E_kin: 97.65376, E_tot: -2.63031, Temperature: 372.85867, phi: -2.7829235, psi: 2.7675953\n",
      "[MindSPONGE] Step: 550, E_pot: -86.63975, E_kin: 79.21601, E_tot: -7.4237366, Temperature: 302.4602, phi: -2.6951704, psi: 2.7429037\n",
      "[MindSPONGE] Step: 560, E_pot: -91.77819, E_kin: 87.55752, E_tot: -4.2206726, Temperature: 334.3095, phi: -2.7110114, psi: 2.6962094\n",
      "[MindSPONGE] Step: 570, E_pot: -82.757126, E_kin: 67.92993, E_tot: -14.827194, Temperature: 259.36804, phi: -2.749873, psi: 2.6291046\n",
      "[MindSPONGE] Step: 580, E_pot: -83.22967, E_kin: 85.018585, E_tot: 1.7889175, Temperature: 324.61542, phi: -2.6595304, psi: 2.6026397\n",
      "[MindSPONGE] Step: 590, E_pot: -78.62119, E_kin: 64.4655, E_tot: -14.155693, Temperature: 246.14024, phi: -2.6389928, psi: 2.6436753\n",
      "[MindSPONGE] Step: 600, E_pot: -105.17256, E_kin: 82.21563, E_tot: -22.956932, Temperature: 313.91324, phi: -2.7518919, psi: 2.6869934\n",
      "[MindSPONGE] Step: 610, E_pot: -87.879456, E_kin: 62.962128, E_tot: -24.917328, Temperature: 240.4001, phi: -2.921755, psi: 2.7449083\n",
      "[MindSPONGE] Step: 620, E_pot: -79.041595, E_kin: 65.41389, E_tot: -13.627708, Temperature: 249.76135, phi: -2.9411612, psi: 2.7441943\n",
      "[MindSPONGE] Step: 630, E_pot: -100.5149, E_kin: 89.814835, E_tot: -10.700066, Temperature: 342.9283, phi: -2.838797, psi: 2.7201588\n",
      "[MindSPONGE] Step: 640, E_pot: -81.88973, E_kin: 74.80992, E_tot: -7.079811, Temperature: 285.63702, phi: -2.7752056, psi: 2.726075\n",
      "[MindSPONGE] Step: 650, E_pot: -90.56945, E_kin: 70.75978, E_tot: -19.80967, Temperature: 270.17288, phi: -2.7366598, psi: 2.7716985\n",
      "[MindSPONGE] Step: 660, E_pot: -99.36987, E_kin: 73.31488, E_tot: -26.054993, Temperature: 279.92868, phi: -2.7629483, psi: 2.845154\n",
      "[MindSPONGE] Step: 670, E_pot: -112.53509, E_kin: 89.88446, E_tot: -22.650627, Temperature: 343.19415, phi: -2.762007, psi: 2.9145272\n",
      "[MindSPONGE] Step: 680, E_pot: -103.703476, E_kin: 74.78784, E_tot: -28.915634, Temperature: 285.5527, phi: -2.8188245, psi: 2.9471686\n",
      "[MindSPONGE] Step: 690, E_pot: -96.31388, E_kin: 57.437744, E_tot: -38.876137, Temperature: 219.30707, phi: -2.791091, psi: 3.0009851\n",
      "[MindSPONGE] Step: 700, E_pot: -103.75195, E_kin: 57.006016, E_tot: -46.745937, Temperature: 217.65866, phi: -2.726889, psi: 3.0908954\n",
      "[MindSPONGE] Step: 710, E_pot: -96.79775, E_kin: 57.007477, E_tot: -39.790276, Temperature: 217.66425, phi: -2.6779416, psi: -3.105036\n",
      "[MindSPONGE] Step: 720, E_pot: -95.55297, E_kin: 76.249214, E_tot: -19.303757, Temperature: 291.13248, phi: -2.6322863, psi: -3.0813298\n",
      "[MindSPONGE] Step: 730, E_pot: -93.776825, E_kin: 84.275345, E_tot: -9.50148, Temperature: 321.77762, phi: -2.586197, psi: 3.1095083\n",
      "[MindSPONGE] Step: 740, E_pot: -116.47222, E_kin: 90.78694, E_tot: -25.68528, Temperature: 346.63998, phi: -2.5499911, psi: 3.0731816\n",
      "[MindSPONGE] Step: 750, E_pot: -84.75309, E_kin: 54.028946, E_tot: -30.724144, Temperature: 206.2917, phi: -2.5638168, psi: 3.1215575\n",
      "[MindSPONGE] Step: 760, E_pot: -114.29947, E_kin: 71.28011, E_tot: -43.019356, Temperature: 272.1596, phi: -2.5500364, psi: 3.1213272\n",
      "[MindSPONGE] Step: 770, E_pot: -100.57666, E_kin: 61.78103, E_tot: -38.79563, Temperature: 235.89047, phi: -2.4918978, psi: 3.05876\n",
      "[MindSPONGE] Step: 780, E_pot: -114.38683, E_kin: 70.71126, E_tot: -43.675575, Temperature: 269.9876, phi: -2.5338905, psi: 2.9935338\n",
      "[MindSPONGE] Step: 790, E_pot: -114.67908, E_kin: 67.86442, E_tot: -46.81466, Temperature: 259.1179, phi: -2.648488, psi: 2.9555407\n",
      "[MindSPONGE] Step: 800, E_pot: -114.864204, E_kin: 61.10498, E_tot: -53.759224, Temperature: 233.30922, phi: -2.7184832, psi: 2.926275\n",
      "[MindSPONGE] Step: 810, E_pot: -96.486885, E_kin: 52.672478, E_tot: -43.814407, Temperature: 201.11247, phi: -2.6822734, psi: 2.9552326\n",
      "[MindSPONGE] Step: 820, E_pot: -122.92024, E_kin: 77.32635, E_tot: -45.593895, Temperature: 295.24515, phi: -2.6272783, psi: 3.0310724\n",
      "[MindSPONGE] Step: 830, E_pot: -93.67839, E_kin: 55.231033, E_tot: -38.447357, Temperature: 210.88148, phi: -2.591601, psi: 3.072984\n",
      "[MindSPONGE] Step: 840, E_pot: -110.61321, E_kin: 67.539024, E_tot: -43.07419, Temperature: 257.87546, phi: -2.4975197, psi: 3.030739\n",
      "[MindSPONGE] Step: 850, E_pot: -111.36262, E_kin: 75.514496, E_tot: -35.84812, Temperature: 288.3272, phi: -2.448619, psi: 2.8941855\n",
      "[MindSPONGE] Step: 860, E_pot: -104.4224, E_kin: 70.876, E_tot: -33.546402, Temperature: 270.6166, phi: -2.471091, psi: 2.7774205\n",
      "[MindSPONGE] Step: 870, E_pot: -77.72243, E_kin: 54.572506, E_tot: -23.149921, Temperature: 208.36711, phi: -2.496534, psi: 2.735218\n",
      "[MindSPONGE] Step: 880, E_pot: -105.01342, E_kin: 71.791985, E_tot: -33.221436, Temperature: 274.114, phi: -2.4535236, psi: 2.809913\n",
      "[MindSPONGE] Step: 890, E_pot: -89.99442, E_kin: 55.76288, E_tot: -34.231544, Temperature: 212.91216, phi: -2.401369, psi: 2.8843958\n",
      "[MindSPONGE] Step: 900, E_pot: -103.018364, E_kin: 60.297146, E_tot: -42.72122, Temperature: 230.22476, phi: -2.4359577, psi: 2.918611\n",
      "[MindSPONGE] Step: 910, E_pot: -103.67759, E_kin: 67.51664, E_tot: -36.16095, Temperature: 257.79, phi: -2.426817, psi: 2.8957512\n",
      "[MindSPONGE] Step: 920, E_pot: -127.50715, E_kin: 86.52681, E_tot: -40.98034, Temperature: 330.37408, phi: -2.4170308, psi: 2.9076355\n",
      "[MindSPONGE] Step: 930, E_pot: -99.001205, E_kin: 69.262375, E_tot: -29.73883, Temperature: 264.45554, phi: -2.4098701, psi: 2.9246547\n",
      "[MindSPONGE] Step: 940, E_pot: -95.82913, E_kin: 64.253136, E_tot: -31.575996, Temperature: 245.32939, phi: -2.4642987, psi: 2.9296904\n",
      "[MindSPONGE] Step: 950, E_pot: -96.32234, E_kin: 70.09251, E_tot: -26.229836, Temperature: 267.62512, phi: -2.6093311, psi: 2.9111068\n",
      "[MindSPONGE] Step: 960, E_pot: -113.833664, E_kin: 77.7847, E_tot: -36.048965, Temperature: 296.9952, phi: -2.6279483, psi: 2.8338194\n",
      "[MindSPONGE] Step: 970, E_pot: -109.855896, E_kin: 79.75572, E_tot: -30.100174, Temperature: 304.5209, phi: -2.6101782, psi: 2.7608547\n",
      "[MindSPONGE] Step: 980, E_pot: -116.46059, E_kin: 95.53732, E_tot: -20.923264, Temperature: 364.77774, phi: -2.59008, psi: 2.7326782\n",
      "[MindSPONGE] Step: 990, E_pot: -101.3766, E_kin: 75.969, E_tot: -25.4076, Temperature: 290.06256, phi: -2.5424967, psi: 2.766436\n",
      "[MindSPONGE] Finished simulation at 2023-05-19 10:24:08\n",
      "[MindSPONGE] Simulation time: 15.86 seconds.\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<mindsponge.core.sponge.Sponge at 0x7fe33b7e2bb0>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cb_h5md = WriteH5MD(system, 'tutorial_a03.h5md', save_freq=10)\n",
    "md.run(1000, callbacks=[run_info, cb_h5md])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mindspore-1.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
